#!/usr/bin/python3

import configparser
import datetime
import subprocess as run
import argparse
import os

# more suitable configparser
class lconfigp(configparser.ConfigParser):
    def optionxform(self, option_str):
        return option_str

# docker address
addresses = {
    "hpl-21.4": "docker://nvcr.io/nvidia/hpc-benchmarks:21.4-hpl",
    "hpcg-21.4": "docker://nvcr.io/nvidia/hpc-benchmarks:21.4-hpcg",
    "hpl-23.3": "docker://nvcr.io/nvidia/hpc-benchmarks:23.3",
    "hpcg-23.3": "docker://nvcr.io/nvidia/hpc-benchmarks:23.3"
}

# hpl.dat
hpl_dat = {
    "device_out": "device out (6=stdout,7=stderr,file)",
    "Ns_counts": "# of problems sizes (N)",
    "Ns": "Ns",
    "NBs_counts": "# of NBs",
    "NBs": "NBs",
    "PMAP_process_mapping": "PMAP process mapping (0=Row-,1=Column-major)",
    "process_grids_counts": "# of process grids (P x Q)",
    "process_grids_P": "Ps",
    "process_grids_Q": "Qs",
    "threshold": "threshold",
    "panel_fact_counts": "# of panel fact",
    "panel_fact_PFACTs": "PFACTs (0=left, 1=Crout, 2=Right)",
    "recursive_stopping_criterium_counts": "# of recursive stopping criterium",
    "recursive_stopping_criterium_NBMINs": "NBMINs (>= 1)",
    "panels_in_recursion": "# of panels in recursion",
    "panels_in_recursion_NDIVs": "NDIVs",
    "recursive_panel_fact_counts": "# of recursive panel fact.",
    "recursive_panel_fact_RFACTs": "RFACTs (0=left, 1=Crout, 2=Right)",
    "broadcast_counts": "# of broadcast",
    "broadcast_BCASTs": "BCASTs (0=1rg,1=1rM,2=2rg,3=2rM,4=Lng,5=LnM)",
    "lookahead_depth_counts": "# of lookahead depth",
    "lookahead_depth_DEPTHs": "DEPTHs (>=0)",
    "SWAP": "SWAP (0=bin-exch,1=long,2=mix)",
    "swapping_threshold": "swapping threshold",
    "L1_in": "L1 in (0=transposed,1=no-transposed) form",
    "U_in": "U  in (0=transposed,1=no-transposed) form",
    "equilibration": "Equilibration (0=no,1=yes)",
    "memory_alignment": "memory alignment in double (> 0)"
}

# hpcg.dat
hpcg_dat = {
    "dimensions": "",
    "seconds": ""
}

# args setting
argp = argparse.ArgumentParser(
    description="RUN Linpack use NVIDIA HPC-Benchmarks",
)
argp.add_argument('--install', type=bool, help='install containers when coressponding conrtainer is not found', const= True, default=False, nargs='?')
# default is log
argp.add_argument('--logdir', type=str, help='specify log directory', default="log")
argp.add_argument('--dryrun', type=bool, help='whether dry run', const=True, default=False, nargs='?')
args = argp.parse_args()
sif_install = args.install
log_dir = os.path.abspath(args.logdir)
dry_run = args.dryrun

# fetch whether singularity in path 
singularity, err = run.Popen(['which', 'singularity'], stdout=run.PIPE, stderr=run.PIPE).communicate()
singularity = singularity.decode('utf-8').strip()

## check singularity and use its abs path
if not os.path.exists(singularity):
    print("singularity is not in your path, run aborted!")
    exit(code=-1)
else:
    singularity = os.path.abspath(singularity)

# check whether logdir exists
if not os.path.exists(args.logdir):
    print("logdir {} does not exist, run aborted!".format(args.logdir))
    exit(code=-1)

# fetch current folder
pwd = os.getcwd()

# fetch configs
con = lconfigp(allow_no_value=True)
con.read("config.ini")

apps = dict(con.items("apps"))
app = apps["app"]
version = apps["version"]
machine = dict(con.items("machine"))
mpi = dict(con.items("mpi"))
numactl_params = dict(con.items("numactl-params"))
hpl = dict(con.items("hpl"))
hpl_numabinds = dict(con.items("hpl-numabinds-{}".format(version)))
hpcg = dict(con.items("hpcg"))
hpcg_numabinds = dict(con.items("hpcg-numabinds-{}".format(version)))

# get params
isMulti = False
mpiparams = ""
numaparams = ""
name = "{}-{}".format(app, version)
if app != "hpl" and app != "hpcg":
    print("no available app found, found {}, run aborted!".format(app))
time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
numabindfile = "{}/numabinds/{}-{}.sh".format(pwd, name, time)
container = "{}/containers/{}.sif".format(pwd, name)
dat = "{}/dats/{}-{}.dat".format(pwd, name, time)
tuning = ""
if app == "hpl":
    tuning = "HPLinpack benchmark input file\nInnovative Computing Laboratory, University of Tennessee\nHPL.out      output file name (if any)\n"
else:
    tuning = "HPCG benchmark input file\nSandia National Laboratories; University of Tennessee, Knoxville\n"
export = ""

if machine["multinodes"] == "yes":
    isMulti = True
    mpiparams += "-H {} ".format(machine["nodes"])

for key, values in mpi.items():
    values = values.strip('\n')
    for value in values.split('\n'):
        mpiparams += "--{} {} ".format(key, value)

for key, value in numactl_params.items():
    numaparams += "--{} {} ".format(key, value)

if app == "hpl":    
    for key, value in hpl.items():
        tuning += "{:<20} {:>30}\n".format(value, hpl_dat[key])
    for key, value in hpl_numabinds.items():
        export += "export {}={}\n".format(key, value)
else:
    for key, value in hpcg.items():
        tuning += "{:<20} {:>30}\n".format(value, hpcg_dat[key])    
    for key, value in hpcg_numabinds.items():
        export += "export {}={}\n".format(key, value)

# check container file exists
if not os.path.isfile(container):
    if sif_install:
        print("container not found, will install it")
        addr = ""
        try:
            addr = addresses.get(name)
        except KeyError:
            print("there is no linpack app or version matched {}", name)
            exit(code = -1)
        os.makedirs(os.path.dirname(container), exist_ok=True)
        print("Now install, you should enter the auth")
        run.Popen(['singularity', 'pull', '--docker-login', container, addr]).communicate()
    else:
        print("container not found, and install is not set, run aborted!")
        exit(code = -1)

# edit params
content = ""
with open("{}/examples/numabinds/{}.sh".format(pwd, name), "r") as file:
    for line in file:
        content += line
if name == "hpl-21.4":
    content = "#!/bin/bash\n{}\n{}".format(export, content)
elif name == "hpl-23.3":
    content = """
#!/bin/bash
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export PMIX_INSTALL_PREFIX=/opt/hpcx/ompi
. /opt/hpcx/hpcx-init.sh
hpcx_load
{}
{}
    """.format(export, content)
elif name == "hpcg-21.4":
    content = "#!/bin/bash\n{}\n{}".format(export, content)
elif name == "hpcg-23.3":
    content = """
#!/bin/bash
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export PMIX_INSTALL_PREFIX=/opt/hpcx/ompi
. /opt/hpcx/hpcx-init.sh
hpcx_load
{}
{}
    """.format(export, content)
else:
    print("there is no availble linpack app or version matched {}", name)
    exit(code = -1)

with open(numabindfile, "w") as file:
    file.write(content)

with open(dat, "w") as file:
    file.write(tuning)

os.chmod(numabindfile, 0o775)
os.chmod(dat, 0o775)

# gen command
command = "mpirun {} {} run -B {} --nv {} {} {} --dat {} 2>&1 | tee {}/{}-{}.log".format(mpiparams, singularity, pwd, container, numabindfile, numaparams, dat, log_dir, name, time)

print("command: {}".format(command))

with open("{}/{}-{}.run".format(log_dir, name, time), "w") as file:
    file.write(command)

if not dry_run:
    run.run(command, shell=True)
